from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.chrome.options import Options as ChromeOptions
from selenium.webdriver.firefox.options import Options as FirefoxOptions
from selenium.webdriver.chrome.service import Service as ChromeService
from selenium.webdriver.firefox.service import Service as FirefoxService
from webdriver_manager.chrome import ChromeDriverManager
from webdriver_manager.firefox import GeckoDriverManager
from selenium.common.exceptions import TimeoutException, WebDriverException
from typing import Dict, List, Optional, Any, Union
import logging
import time
from crawler_base import CrawlerBase

logger = logging.getLogger(__name__)

class SeleniumCrawler(CrawlerBase):
    def __init__(self, config=None):
        super().__init__(config)
        self.driver = None
        self.wait = None
        self._setup_driver()
    
    def _setup_driver(self):
        selenium_config = self.config.get('selenium_settings', {})
        driver_type = selenium_config.get('driver', 'chrome').lower()
        headless = selenium_config.get('headless', True)
        window_size = selenium_config.get('window_size', [1920, 1080])
        page_load_timeout = selenium_config.get('page_load_timeout', 30)
        
        try:
            if driver_type == 'chrome':
                self.driver = self._create_chrome_driver(headless, window_size)
            elif driver_type == 'firefox':
                self.driver = self._create_firefox_driver(headless, window_size)
            else:
                raise ValueError(f"Unsupported driver type: {driver_type}")
            
            self.driver.set_page_load_timeout(page_load_timeout)
            self.wait = WebDriverWait(self.driver, 10)
            
            logger.info(f"Selenium {driver_type} driver initialized successfully")
            
        except Exception as e:
            logger.error(f"Failed to initialize Selenium driver: {str(e)}")
            raise
    
    def _create_chrome_driver(self, headless: bool, window_size: List[int]) -> webdriver.Chrome:
        options = ChromeOptions()
        
        if headless:
            options.add_argument('--headless')
        
        options.add_argument(f'--window-size={window_size[0]},{window_size[1]}')
        options.add_argument('--no-sandbox')
        options.add_argument('--disable-dev-shm-usage')
        options.add_argument('--disable-gpu')
        options.add_argument('--disable-extensions')
        options.add_argument('--disable-images')
        options.add_argument('--disable-javascript')
        options.add_argument(f'--user-agent={self.get_random_user_agent()}')
        
        # 设置代理
        proxy = self.get_next_proxy()
        if proxy and 'http' in proxy:
            proxy_address = proxy['http'].replace('http://', '')
            options.add_argument(f'--proxy-server={proxy_address}')
        
        service = ChromeService(ChromeDriverManager().install())
        return webdriver.Chrome(service=service, options=options)
    
    def _create_firefox_driver(self, headless: bool, window_size: List[int]) -> webdriver.Firefox:
        options = FirefoxOptions()
        
        if headless:
            options.add_argument('--headless')
        
        options.set_preference("general.useragent.override", self.get_random_user_agent())
        
        # 设置代理
        proxy = self.get_next_proxy()
        if proxy and 'http' in proxy:
            proxy_parts = proxy['http'].replace('http://', '').split(':')
            if len(proxy_parts) == 2:
                options.set_preference("network.proxy.type", 1)
                options.set_preference("network.proxy.http", proxy_parts[0])
                options.set_preference("network.proxy.http_port", int(proxy_parts[1]))
        
        service = FirefoxService(GeckoDriverManager().install())
        driver = webdriver.Firefox(service=service, options=options)
        driver.set_window_size(window_size[0], window_size[1])
        return driver
    
    def get_page(self, url: str, wait_element: Optional[str] = None, wait_time: int = 10, **kwargs) -> Optional[str]:
        try:
            logger.info(f"Loading page: {url}")
            self.driver.get(url)
            
            # 等待特定元素加载
            if wait_element:
                try:
                    element_wait = WebDriverWait(self.driver, wait_time)
                    if wait_element.startswith('//'):
                        element_wait.until(EC.presence_of_element_located((By.XPATH, wait_element)))
                    else:
                        element_wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, wait_element)))
                    logger.info(f"Wait element found: {wait_element}")
                except TimeoutException:
                    logger.warning(f"Wait element not found within {wait_time}s: {wait_element}")
            
            # 执行JavaScript（如果提供）
            if 'javascript' in kwargs:
                self.driver.execute_script(kwargs['javascript'])
                time.sleep(1)
            
            self.add_delay()
            return self.driver.page_source
            
        except WebDriverException as e:
            logger.error(f"Selenium error loading page {url}: {str(e)}")
            return None
        except Exception as e:
            logger.error(f"Unexpected error: {str(e)}")
            return None
    
    def click_element(self, selector: str, by_type: str = 'css') -> bool:
        try:
            if by_type.lower() == 'xpath':
                element = self.wait.until(EC.element_to_be_clickable((By.XPATH, selector)))
            else:
                element = self.wait.until(EC.element_to_be_clickable((By.CSS_SELECTOR, selector)))
            
            element.click()
            logger.info(f"Clicked element: {selector}")
            self.add_delay()
            return True
            
        except Exception as e:
            logger.error(f"Failed to click element {selector}: {str(e)}")
            return False
    
    def fill_form(self, form_data: Dict[str, str], by_type: str = 'css') -> bool:
        try:
            for selector, value in form_data.items():
                if by_type.lower() == 'xpath':
                    element = self.wait.until(EC.presence_of_element_located((By.XPATH, selector)))
                else:
                    element = self.wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, selector)))
                
                element.clear()
                element.send_keys(value)
                logger.info(f"Filled form field {selector} with value")
            
            self.add_delay()
            return True
            
        except Exception as e:
            logger.error(f"Failed to fill form: {str(e)}")
            return False
    
    def scroll_page(self, direction: str = 'down', pixels: int = 1000) -> bool:
        try:
            if direction.lower() == 'down':
                self.driver.execute_script(f"window.scrollBy(0, {pixels});")
            elif direction.lower() == 'up':
                self.driver.execute_script(f"window.scrollBy(0, -{pixels});")
            elif direction.lower() == 'bottom':
                self.driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
            elif direction.lower() == 'top':
                self.driver.execute_script("window.scrollTo(0, 0);")
            
            time.sleep(1)
            return True
            
        except Exception as e:
            logger.error(f"Failed to scroll page: {str(e)}")
            return False
    
    def wait_for_element(self, selector: str, by_type: str = 'css', timeout: int = 10) -> bool:
        try:
            wait = WebDriverWait(self.driver, timeout)
            if by_type.lower() == 'xpath':
                wait.until(EC.presence_of_element_located((By.XPATH, selector)))
            else:
                wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, selector)))
            
            logger.info(f"Element found: {selector}")
            return True
            
        except TimeoutException:
            logger.warning(f"Element not found within {timeout}s: {selector}")
            return False
    
    def take_screenshot(self, filepath: str) -> bool:
        try:
            self.driver.save_screenshot(filepath)
            logger.info(f"Screenshot saved: {filepath}")
            return True
        except Exception as e:
            logger.error(f"Failed to take screenshot: {str(e)}")
            return False
    
    def get_element_text(self, selector: str, by_type: str = 'css') -> Optional[str]:
        try:
            if by_type.lower() == 'xpath':
                element = self.driver.find_element(By.XPATH, selector)
            else:
                element = self.driver.find_element(By.CSS_SELECTOR, selector)
            
            return element.text
            
        except Exception as e:
            logger.error(f"Failed to get element text {selector}: {str(e)}")
            return None
    
    def get_elements_text(self, selector: str, by_type: str = 'css') -> List[str]:
        try:
            if by_type.lower() == 'xpath':
                elements = self.driver.find_elements(By.XPATH, selector)
            else:
                elements = self.driver.find_elements(By.CSS_SELECTOR, selector)
            
            return [elem.text for elem in elements if elem.text]
            
        except Exception as e:
            logger.error(f"Failed to get elements text {selector}: {str(e)}")
            return []
    
    def close(self):
        if self.driver:
            try:
                self.driver.quit()
                logger.info("Selenium driver closed")
            except Exception as e:
                logger.error(f"Error closing driver: {str(e)}")